package com.zhora.ai.v1.business.impl;

import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhora.ai.config.minio.MinIoConfig;
import com.zhora.ai.v1.dto.AiLibraryDocDetailDTO;
import com.zhora.ai.v1.dto.AiLibraryDocListDTO;
import com.zhora.ai.v1.dto.AiLibraryDocSegmentListDTO;
import com.zhora.common.utils.ConvertUtils;
import com.zhora.dto.ai.AiLibraryDocDTO;
import com.zhora.dto.ai.AiLibraryDocSegmentDTO;
import com.zhora.dto.ai.search.AiLibraryDocSearchDTO;
import com.zhora.dto.ai.search.AiLibraryDocSegmentSearchDTO;
import com.zhora.dto.system.SecurityUser;
import com.zhora.entity.ai.AiLibraryDocSegmentEntity;
import com.zhora.enums.ai.Actions;
import com.zhora.enums.ai.LibraryDocStatusEnum;
import com.zhora.service.ai.IAiLibraryDocSegmentService;
import com.zhora.service.ai.IAiLibraryDocService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.amazon.s3.AmazonS3DocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.RandomStringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Configuration
@RequiredArgsConstructor
@Service
@Slf4j
public class RagService {

    @Autowired
    private EmbeddingModel qwenEmbeddingModel;

    @Autowired
    @Qualifier("qwenEmbeddingStore")
    private EmbeddingStore<TextSegment> qwenEmbeddingStore;

    @Autowired
    private IAiLibraryDocService libraryDocService;

    private final MinIoConfig minIoConfig;

    private final AmazonS3DocumentLoader amazonS3DocumentLoader;
    private final EmbeddingModel aliEmbeddingModel;
    private final EmbeddingStore<TextSegment> aliEmbeddingStore;
    private final IAiLibraryDocService aiLibraryDocService;
    private final IAiLibraryDocSegmentService aiLibraryDocSegmentService;

    public void deleteLibraryBy(Long libraryId) {
        AiLibraryDocSearchDTO docSearchDTO = new AiLibraryDocSearchDTO();
        docSearchDTO.setLibId(libraryId);
        docSearchDTO.setDelFlag(Boolean.FALSE);
        List<AiLibraryDocListDTO> libraryDocs = ConvertUtils.sourceToTarget(libraryDocService.list(docSearchDTO), AiLibraryDocListDTO.class);
        List<Long> docIds = libraryDocs.stream().map(AiLibraryDocListDTO::getId).toList();
        for (Long docId : docIds) {
            deleteDocBy(docId);
        }
        aiLibraryDocService.removeByLibraryId(libraryId);
    }

    public void deleteDocBy(Long docId) {
        AiLibraryDocSegmentSearchDTO segmentSearchDTO = new AiLibraryDocSegmentSearchDTO();
        segmentSearchDTO.setDocId(docId);
        segmentSearchDTO.setDelFlag(Boolean.FALSE);

        List<AiLibraryDocSegmentListDTO> libraryDocSegments = ConvertUtils.sourceToTarget(aiLibraryDocSegmentService.list(segmentSearchDTO), AiLibraryDocSegmentListDTO.class);
        List<String> embeddingIdList =
                libraryDocSegments.stream().map(AiLibraryDocSegmentListDTO::getEmbeddingId).toList();
        if (CollectionUtils.isNotEmpty(embeddingIdList)) {
            aliEmbeddingStore.removeAll(embeddingIdList);
        }
        aiLibraryDocService.removeById(docId);
    }

    public Map<String, String> searchAction(String message) {
        Map<String, String> result = new HashMap<>();
        EmbeddingSearchRequest embeddingSearchRequest =
                EmbeddingSearchRequest.builder()
                        .queryEmbedding(qwenEmbeddingModel.embed(message).content())
                        .minScore(0.89)
                        .build();
        EmbeddingSearchResult<TextSegment> embeddingSearchResult = qwenEmbeddingStore.search(embeddingSearchRequest);
        if (!embeddingSearchResult.matches().isEmpty()) {
            Metadata metadata = embeddingSearchResult.matches().getFirst().embedded().metadata();
            result.put(Actions.INDEX_KEY, metadata.getString(Actions.INDEX_KEY));
        }
        return result;
    }

    public Long createLibraryDocBy(Long libraryId, String objectName, String originalName)
            throws JsonProcessingException {
        String username = SecurityUser.getUser().getUsername();
        String identify =
                String.format(
                        "%d%s_%s",
                        Instant.now().toEpochMilli(),
                        RandomStringUtils.insecure().nextAlphabetic(6),
                        originalName);
        Map<String, String> meta = new HashMap<>();
        meta.put("uploader", username);
        AiLibraryDocDTO libraryDoc = new AiLibraryDocDTO();
        ObjectMapper objectMapper = new ObjectMapper();
        String metaJson = objectMapper.writeValueAsString(meta);
        libraryDoc.setMeta(ConvertUtils.sourceToTarget(JSONUtil.parse(metaJson), Map.class));
        libraryDoc.setPath(objectName);
        libraryDoc.setName(originalName);
        libraryDoc.setIdentify(identify);
        libraryDoc.setLibId(libraryId);
        libraryDoc.setStatus(LibraryDocStatusEnum.INDEXING.getValue());
        libraryDoc.setEnable(Boolean.TRUE);
        libraryDocService.create(libraryDoc);

        AiLibraryDocSearchDTO libraryDocSearchDTO = new AiLibraryDocSearchDTO();
        libraryDocSearchDTO.setIdentify(identify);
        libraryDocSearchDTO.setDelFlag(Boolean.FALSE);

        return libraryDocService.getDetail(libraryDocSearchDTO).getId();
    }

    @Async
    public void embeddingAndCreateDocSegment(Long libraryId, Long libraryDocId, String objectName) {
        Document document =
                amazonS3DocumentLoader.loadDocument(
                        minIoConfig.getDefaultBucket(), objectName, new ApacheTikaDocumentParser());
        List<AiLibraryDocSegmentDTO> libraryDocSegmentList = new ArrayList<>();
        DocumentByParagraphSplitter documentByParagraphSplitter =
                new DocumentByParagraphSplitter(500, 150);

        documentByParagraphSplitter
                .split(document)
                .forEach(
                        textSegment -> {
                            Response<Embedding> embed = aliEmbeddingModel.embed(textSegment);
                            Integer tokenUsage = embed.tokenUsage().totalTokenCount();
                            Embedding vector = embed.content();
                            textSegment.metadata().put("libraryId", libraryId);
                            String embeddingId = aliEmbeddingStore.add(vector, textSegment);
                            AiLibraryDocSegmentDTO libraryDocSegment = new AiLibraryDocSegmentDTO();
                            libraryDocSegment.setEmbeddingId(embeddingId);
                            libraryDocSegment.setContent(textSegment.text());
                            libraryDocSegment.setTokenUsage(tokenUsage);
                            libraryDocSegment.setDocId(libraryDocId);
                            libraryDocSegmentList.add(libraryDocSegment);
                        });
        aiLibraryDocSegmentService.saveBatch(ConvertUtils.sourceToTarget(libraryDocSegmentList, AiLibraryDocSegmentEntity.class));

        AiLibraryDocSearchDTO libraryDocSearchDTO = new AiLibraryDocSearchDTO();
        libraryDocSearchDTO.setId(libraryDocId);
        libraryDocSearchDTO.setDelFlag(Boolean.FALSE);

        AiLibraryDocDTO libraryDoc = aiLibraryDocService.getDetail(libraryDocSearchDTO);
        libraryDoc.setStatus(LibraryDocStatusEnum.SUCCESS.getValue());
        aiLibraryDocService.updateById(libraryDoc);
    }
}
